package com.wt.admin.service.image.impl;

import cn.hutool.core.bean.BeanUtil;
import cn.hutool.core.collection.CollUtil;
import cn.hutool.core.io.FileUtil;
import cn.hutool.core.util.ObjectUtil;
import cn.hutool.json.JSONUtil;
import com.wt.admin.config.ConstVar;
import com.wt.admin.config.prop.IndexProp;
import com.wt.admin.config.socket.WebSocketSessionManager;
import com.wt.admin.domain.dto.image.ImageClassificationDTO;
import com.wt.admin.domain.entity.image.ImageItemEntity;
import com.wt.admin.domain.entity.model.ModelListEntity;
import com.wt.admin.domain.model.ImageYoloModel;
import com.wt.admin.domain.vo.socket.ProgressVO;
import com.wt.admin.domain.vo.socket.SocketVO;
import com.wt.admin.util.FileUtils;
import jakarta.annotation.Resource;
import lombok.SneakyThrows;
import lombok.extern.slf4j.Slf4j;
import org.dromara.easyai.entity.ThreeChannelMatrix;
import org.dromara.easyai.tools.ImageTools;
import org.dromara.easyai.tools.Picture;
import org.dromara.easyai.yolo.*;
import org.springframework.stereotype.Service;

import java.io.File;
import java.io.InputStream;
import java.util.ArrayList;
import java.util.List;
import java.util.function.Consumer;
import java.util.function.Function;

@Service
@Slf4j
public class ImageTrainingService {

    @Resource
    private IndexProp indexProp;

    @SneakyThrows
    public void training(ImageClassificationDTO data, List<ImageItemEntity> imageItemVOList,String modelName, Function<Boolean, ModelListEntity> fun) {
        FastYolo fastYolo = new FastYolo(data.getYoloConfig());
        List<YoloSample> yoloSamples = new ArrayList<>();
        imageItemVOList.forEach(image -> {
            String imageAbsURL = FileUtils.getImageAbsURL(image.getImageUrl());
            YoloSample yoloSample = new YoloSample();
            yoloSample.setLocationURL(imageAbsURL);
            List<YoloBody> yoloBodies = new ArrayList<>();
            List<Object> featuresConfig = image.getFeaturesConfig();
            if(ObjectUtil.isEmpty(featuresConfig)){
                return;
            }
            featuresConfig.forEach(i -> {
                YoloBody yoloBody = new YoloBody();
                BeanUtil.copyProperties(i, yoloBody);
                yoloBodies.add(yoloBody);
            });
            yoloSample.setYoloBodies(yoloBodies);
            yoloSamples.add(yoloSample);
        });
        fastYolo.toStudy(yoloSamples,null);
        YoloModel model = fastYolo.getModel();
        ImageYoloModel imageYoloModel = new ImageYoloModel();
        imageYoloModel.setYoloConfig(data.getYoloConfig());
        imageYoloModel.setYoloModel(model);
        imageYoloModel.setImageWidth(data.getImageWidth());
        imageYoloModel.setImageHeight(data.getImageHeight());
        ModelListEntity entity = fun.apply(true);
        FileUtil.writeUtf8String(JSONUtil.toJsonStr(imageYoloModel), String.format(indexProp.getModelPath().getYoloModel(), modelName));
    }


    @SneakyThrows
    public List<OutBox> testTraining(String model,String fromUrl,String toUrl) {
        FastYolo fastYolo = toFastYolo(model);
        ThreeChannelMatrix matrix = Picture.getThreeMatrix(fromUrl, false);
        List<OutBox> look = fastYolo.look(matrix, System.currentTimeMillis());
        if(CollUtil.isEmpty(look)){
            look = new ArrayList<>();
        }
        // 验证，把识别内容写到图片中
        ImageTools.drawBox(fromUrl,look, toUrl,14);
        return look;
    }

    @SneakyThrows
    public List<OutBox> testTraining(String model, File file, String toUrl, Consumer<ImageYoloModel> consumer) {
        ImageYoloModel bean = toImageYoloModel(model);
        consumer.accept( bean);
        FastYolo fastYolo = new FastYolo(bean.getYoloConfig());
        fastYolo.insertModel(bean.getYoloModel());
        ThreeChannelMatrix matrix = Picture.getThreeMatrix(file, false);
        List<OutBox> look = fastYolo.look(matrix, System.currentTimeMillis());
        if(CollUtil.isEmpty(look)){
            look = new ArrayList<>();
        }
        // 验证，把识别内容写到图片中
        ImageTools.drawBox(file.getAbsolutePath(),look, toUrl,14);
        return look;
    }

    @SneakyThrows
    public List<OutBox> testTraining(String model,File file,String toUrl) {
        FastYolo fastYolo = toFastYolo(model);
        ThreeChannelMatrix matrix = Picture.getThreeMatrix(file, false);
        List<OutBox> look = fastYolo.look(matrix, System.currentTimeMillis());
        if(CollUtil.isEmpty(look)){
            look = new ArrayList<>();
        }
        // 验证，把识别内容写到图片中
        ImageTools.drawBox(file.getAbsolutePath(),look, toUrl,14);
        return look;
    }

    @SneakyThrows
    public List<OutBox> testTraining(String model,File file) {
        FastYolo fastYolo = toFastYolo(model);
        ThreeChannelMatrix matrix = Picture.getThreeMatrix(file, false);
        List<OutBox> look = fastYolo.look(matrix, System.currentTimeMillis());
        if(CollUtil.isEmpty(look)){
            look = new ArrayList<>();
        }
        return look;
    }

    @SneakyThrows
    public List<OutBox> testTraining(String model, InputStream file) {
        FastYolo fastYolo = toFastYolo(model);
        ThreeChannelMatrix matrix = Picture.getThreeMatrix(file, false);
        List<OutBox> look = fastYolo.look(matrix, System.currentTimeMillis());
        if(CollUtil.isEmpty(look)){
            look = new ArrayList<>();
        }
        return look;
    }

    @SneakyThrows
    public List<OutBox> testTraining(FastYolo fastYolo, InputStream file) {
        ThreeChannelMatrix matrix = Picture.getThreeMatrix(file, false);
        List<OutBox> look = fastYolo.look(matrix, System.currentTimeMillis());
        if(CollUtil.isEmpty(look)){
            look = new ArrayList<>();
        }
        return look;
    }

    @SneakyThrows
    public List<OutBox> testTraining(String model,String path) {
        FastYolo fastYolo = toFastYolo(model);
        ThreeChannelMatrix matrix = Picture.getThreeMatrix(path, false);
        List<OutBox> look = fastYolo.look(matrix, System.currentTimeMillis());
        if(CollUtil.isEmpty(look)){
            look = new ArrayList<>();
        }
        return look;
    }

    public FastYolo toFastYolo(String model) throws Exception {
        File file = FileUtil.file(String.format(indexProp.getModelPath().getYoloModel(),model));
        if(!file.exists()){
            return null;
        }
        String json = FileUtil.readUtf8String(file);
        ImageYoloModel bean = JSONUtil.toBean(json, ImageYoloModel.class);
        FastYolo fastYolo = new FastYolo(bean.getYoloConfig());
        fastYolo.insertModel(bean.getYoloModel());
        return fastYolo;
    }

    public ImageYoloModel toImageYoloModel(String model) throws Exception {
        File file = FileUtil.file(String.format(indexProp.getModelPath().getYoloModel(),model));
        if(!file.exists()){
            return null;
        }
        String json = FileUtil.readUtf8String(file);
        return JSONUtil.toBean(json, ImageYoloModel.class);
    }


    private void send(int current){
        WebSocketSessionManager.sendToAll(new SocketVO<ProgressVO>
                (ConstVar.Socket.PROGRESS,ProgressVO.set(current)));
    }


}
